import warnings

warnings.filterwarnings('ignore')
import numpy as np
import tensorflow as tf
import gym
from gym.spaces import Box
from scipy.misc import imresize
import random
import cv2
import time
import logging
import os
import sys

# We
# define
# the
# Class
# EnvWrapper and define
# some
# of
# the
# environment
# wrapper
# functions
#
# In[2]:


class EnvWrapper:

    # First we define the __init__ method and initialize variables

    def __init__(self, env_name, debug=False):

        # environment name
        self.env_name = env_name

        # initialize the gym environment
        self.env = gym.make(env_name)

        # get the action space
        self.action_space = self.env.action_space

        # get the observation space
        self.observation_space = Box(low=0, high=255, shape=(84, 84, 4))

        # initialize frame_num for storing the frame count
        self.frame_num = 0

        # For recording the game screen
        self.monitor = self.env.monitor

        # initialize frames
        self.frames = np.zeros((84, 84, 4), dtype=np.uint8)

        # initialize a boolean called debug when set true last few frames will be displayed
        self.debug = debug

        if self.debug:
            cv2.startWindowThread()
            cv2.namedWindow("Game")

    # we define the function called step where we perform some action in the
    # environment, receive reward and move to the next state
    # step function will take the current state as input and returns the preprocessed frame as next state

    def step(self, a):
        ob, reward, done, xx = self.env.step(a)
        return self.process_frame(ob), reward, done, xx

    # We define the helper function called reset for resetting the environment
    # after resetting it will return the preprocessed game screen

    def reset(self):
        self.frame_num = 0
        return self.process_frame(self.env.reset())

    # next we define another helper function for rendering the environment
    def render(self):
        return self.env.render()

    # now we define the function called process_frame for preprocessing the frame

    def process_frame(self, frame):

        # convert the image to gray
        state_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # change the size
        state_resized = cv2.resize(state_gray, (84, 110))

        # resize
        gray_final = state_resized[16:100, :]

        if self.frame_num == 0:
            self.frames[:, :, 0] = gray_final
            self.frames[:, :, 1] = gray_final
            self.frames[:, :, 2] = gray_final
            self.frames[:, :, 3] = gray_final

        else:
            self.frames[:, :, 3] = self.frames[:, :, 2]
            self.frames[:, :, 2] = self.frames[:, :, 1]
            self.frames[:, :, 1] = self.frames[:, :, 0]
            self.frames[:, :, 0] = gray_final

        # increment the frame_num counter

        self.frame_num += 1

        if self.debug:
            cv2.imshow('Game', gray_final)

        return self.frames.copy()